import os
import cv2
import khandy
import pandas as pd
import numpy as np
import streamlit as st
from insectid import InsectDetector, InsectIdentifier
import time
import requests
import base64
import io
from PIL import Image

# --- 核心新增：导入Modal组件 ---
from streamlit_modal import Modal

# --- API 和 聊天逻辑函数 (保持不变) ---
API_URL = "http://localhost:8503/chat"
VIT_API_URL = "http://localhost:8503/vit_inference"
def real_multimodal_chatbot(image_np, species_info, user_query, chat_history):
    # ... (此函数代码完全不需要修改)
    try:
        image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB); pil_image = Image.fromarray(image_rgb)
        buffered = io.BytesIO(); pil_image.save(buffered, format="JPEG"); img_str_base64 = base64.b64encode(buffered.getvalue()).decode()
        species_info_for_api = species_info.copy(); 
        if 'image' in species_info_for_api: del species_info_for_api['image']
        history_for_api = [{"role": msg["role"], "content": msg["content"]} for msg in chat_history]
        payload = {"image_base64": img_str_base64, "prompt": user_query, "species_info": species_info_for_api, "chat_history": history_for_api}
        response = requests.post(API_URL, json=payload, timeout=120); response.raise_for_status()
        result = response.json()
        if "response" in result:
            for word in result["response"].split(): yield word + " "; time.sleep(0.02)
        elif "error" in result: yield f"模型API返回错误: {result['error']}"
    except requests.exceptions.RequestException as e:
        yield (f"无法连接到模型API服务，请确认：\n1. 后端API服务 (`api_server.py`) 是否已在服务器上成功启动？\n2. `API_URL` (`{API_URL}`) 是否设置正确？\n详细错误: {e}")
    except Exception as e: yield f"与模型交互时发生未知错误: {e}"

# --- Session State 初始化 ---
def reset_state():
    st.session_state.processed_images = {}
    st.session_state.active_chat_image = None
    st.session_state.selected_files = [] # 存储从模态框选择的文件
    st.session_state["camera_image_ready"] = False
    st.session_state["camera_image_bytes"] = None
    st.session_state["chat_history"] = []
    # st.session_state["current_insect_info"] = None

# 初始化所有需要的session state键
if "chat_history" not in st.session_state:
    st.session_state["chat_history"] = []
if "camera_image_ready" not in st.session_state:
    st.session_state["camera_image_ready"] = False
    st.session_state["camera_image_bytes"] = None
if 'processed_images' not in st.session_state: 
    st.session_state.processed_images = {} 
if 'active_chat_image' not in st.session_state: 
    st.session_state.active_chat_image = None
if 'selected_files' not in st.session_state: 
    st.session_state.selected_files = []

# --- 页面配置和UI ---
st.set_page_config(page_title="昆虫识别演示", layout="wide")
st.title("昆虫识别演示")
st.write("选择图片或通过摄像头拍摄，自动检测并识别物种。可与AI助手针对图片内容进行对话。")
st.markdown("""
    <style>
    /* 调整 st.camera_input 容器的整体大小 */
    .stCameraInput {
        width: 640px;
        height: 650px; /* 增加高度以容纳拍摄按钮 */
        margin-left: auto;
        margin-right: auto;
        overflow: visible; /* 确保按钮不被裁剪 */
    }
    /* 调整视频元素大小 */
    .stCameraInput > div > div > video {
        width: 550px;
        height: 500px;
        # display: block;
        margin-left: auto;
        margin-right: auto;
        margin-top: -400px;
    }
    .stCameraInput button {
        width: 550px !important;
        margin-top: -475px !important;
        display: block;
        margin-left: auto;
        margin-right: auto;
    }
    </style>
""", unsafe_allow_html=True)
st.markdown("""
    <style>
    .stCameraInput .stImage {
        display: none !important;
    }
    .stCameraInput img {
        display: none !important;
    }
    .stCameraInput div[data-testid="stImage"] {
        display: none !important;
    }
    /* 让摄像头组件整体更紧凑 */
    .stCameraInput {
        min-height: 0px !important;
        padding-bottom: 0px !important;
    }
    </style>
""", unsafe_allow_html=True)

# --- 全新的、基于模态窗口的输入处理逻辑 ---

# ⚠️ 定义服务器上存放图片的基础路径
SERVER_IMAGE_BASE_PATH = "/data/users/guochao/datasets/insectid_test_images" 

def load_image_from_server(filepath):
    try: return cv2.imread(filepath)
    except Exception: return None

# --- 创建模态窗口实例 ---
# 我们将复用这个模态窗口来处理单选和多选
file_selector_modal = Modal(
    "从服务器选择图片", 
    key="file_selector_modal",
    max_width=1200 # 设置一个较大的宽度以容纳图片网格
)

# --- 主页面UI ---
mode = st.radio(
    "选择输入方式",
    ["选择图片", "摄像头拍摄"],
    on_change=reset_state,
    key="main_mode_selector"
)

img_files = st.session_state.get('selected_files', [])

if mode == "选择图片":
    col1, col2 = st.columns(2)

    with col1:
        if st.button("选择单张图片进行识别", type="primary"):
            st.session_state.current_path = SERVER_IMAGE_BASE_PATH
            st.session_state.selection_mode = "single"
            file_selector_modal.open()
    with col2:
        if st.button("选择多张图片进行识别", type="primary"):
            st.session_state.current_path = SERVER_IMAGE_BASE_PATH
            st.session_state.selection_mode = "multiple"
            file_selector_modal.open()

elif mode == "摄像头拍摄":
    st.info("请拍摄一张清晰的昆虫照片，确保光线充足且昆虫位于画面中心。")

    if not st.session_state["camera_image_ready"]:
        camera_image = st.camera_input("拍摄照片", on_change=reset_state, key="camera_input_widget")
        # print(1)
        if camera_image:
            st.session_state["camera_image_bytes"] = camera_image.getvalue()
            st.session_state["camera_image_ready"] = True
            st.rerun()
    else:
        file_bytes = np.asarray(bytearray(st.session_state["camera_image_bytes"]), dtype=np.uint8)
        image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
        if image is not None: 
            st.session_state.selected_files = [("摄像头照片",image)]
        if st.button("重新拍摄"):
            reset_state()
            st.rerun()

# --- 模态窗口内容定义 ---
if file_selector_modal.is_open():
    with file_selector_modal.container():
        st.subheader("请选择您要处理的图片")
        current_path = st.session_state.get('current_path', SERVER_IMAGE_BASE_PATH)
        relative_path = os.path.relpath(current_path, SERVER_IMAGE_BASE_PATH)
        
        if relative_path == '.': relative_path = '/' # 根目录显示为'/'
        
        nav_col1, nav_col2 = st.columns([4, 1])
        # with nav_col1:
            # st.subheader(f"当前路径: {relative_path}")
        with nav_col2:
            # 只有不在根目录时才显示“返回”按钮
            if current_path != SERVER_IMAGE_BASE_PATH:
                if st.button("⬅️ 返回上一级", use_container_width=True):
                    st.session_state.current_path = os.path.dirname(current_path)
                    st.rerun()
        
        st.markdown("---")

        # --- 2. 遍历当前路径下的内容 ---
        try:
            items = os.listdir(current_path)
            directories = sorted([d for d in items if os.path.isdir(os.path.join(current_path, d))])
            image_files = sorted([f for f in items if f.lower().endswith(('.jpg', '.jpeg', '.png'))])

            # --- 3. 渲染文件夹列表 ---
            if directories:
                # st.markdown("#### 文件夹")
                cols_per_row = 4
                dir_cols = st.columns(cols_per_row)
                for i, dir_name in enumerate(directories):
                    with dir_cols[i % cols_per_row]:
                        if st.button(f"📁 {dir_name}", use_container_width=True, help=f"进入 {dir_name} 文件夹"):
                            st.session_state.current_path = os.path.join(current_path, dir_name)
                            st.rerun()
                st.markdown("---")
        except Exception as e:
            st.error(f"无法访问路径 {current_path}: {e}")

        all_image_paths = []
        if os.path.isdir(SERVER_IMAGE_BASE_PATH):
            for root, _, files in os.walk(SERVER_IMAGE_BASE_PATH):
                for file in files:
                    if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                        all_image_paths.append(os.path.join(root, file))
        
        if not all_image_paths:
            st.warning("在指定路径下没有找到任何图片文件。")
        else:
            selection_mode = st.session_state.get("selection_mode", "single")
            
            cols_per_row = 4
            cols = st.columns(cols_per_row)
            
            if selection_mode == "single":
                for i, fname in enumerate(image_files):
                    with cols[i % cols_per_row]:
                        full_path = os.path.join(current_path, fname)
                        st.image(full_path, use_container_width=True)
                        if st.button(f"选择这张", key=f"modal_btn_{full_path}", use_container_width=True):
                            reset_state()
                            image = load_image_from_server(full_path)
                            if image is not None:
                                st.session_state.selected_files = [(fname, image)]
                            file_selector_modal.close()
                            st.rerun()


            elif selection_mode == "multiple":
                if 'modal_checkbox_states' not in st.session_state:
                    st.session_state.modal_checkbox_states = {}
                
                for i, fname in enumerate(image_files):
                    with cols[i % cols_per_row]:
                        full_path = os.path.join(current_path, fname)
                        st.image(full_path, use_container_width=True)
                        st.session_state.modal_checkbox_states[full_path] = st.checkbox(
                            "勾选", key=f"modal_cb_{full_path}", value=st.session_state.modal_checkbox_states.get(full_path, False)
                        )
                
                st.markdown("---")
                if st.button("✅ 确认选择", type="primary", use_container_width=True):
                    reset_state()
                    selected_images = []
                    for path, is_selected in st.session_state.modal_checkbox_states.items():
                        if is_selected:
                            image = load_image_from_server(path)
                            if image is not None:
                                selected_images.append((os.path.basename(path), image))
                    st.session_state.selected_files = selected_images
                    file_selector_modal.close()
                    st.rerun()


            # else: # 选择多张图片
            #     # 使用 st.container 来包裹整个多选区域

            #     # --- 1. 创建顶部的固定行 ---
            #     top_col1, top_col2 = st.columns([3, 1]) # 左侧占3份，右侧占1份

            #     # with top_col1:
            #     #     st.subheader("请勾选您想选择的图片，然后点击右侧按钮确认")

            #     with top_col2:
            #         # 将确认按钮放在这里
            #         if st.button("✅ 确认选择", type="primary", use_container_width=True):
            #             # 确认按钮的逻辑
            #             reset_state()
            #             selected_images = []
            #             # 从 session_state 中获取勾选状态
            #             for path, is_selected in st.session_state.get('modal_checkbox_states', {}).items():
            #                 if is_selected:
            #                     image = load_image_from_server(path)
            #                     if image is not None:
            #                         selected_images.append((os.path.basename(path), image))
            #             st.session_state.selected_files = selected_images
            #             st.rerun()

            #     st.markdown("---") # 添加一条分割线

            #     # --- 2. 渲染可滚动的图片网格 ---
            #     cols_per_row = 5
            #     cols = st.columns(cols_per_row)
                
            #     # 使用一个字典来存储每个checkbox的状态
            #     if 'modal_checkbox_states' not in st.session_state:
            #         st.session_state.modal_checkbox_states = {p: False for p in all_image_paths}

            #     for i, full_path in enumerate(all_image_paths):
            #         col = cols[i % cols_per_row]
            #         with col:
            #             st.image(full_path, use_container_width=True, caption=os.path.basename(full_path))
            #             # 为每个checkbox提供唯一的key，并与session_state关联
            #             is_checked = st.checkbox(
            #                 "勾选", 
            #                 key=f"cb_{full_path}", 
            #                 value=st.session_state.modal_checkbox_states.get(full_path, False)
            #             )
            #             # 实时更新session_state中的勾选状态
            #             st.session_state.modal_checkbox_states[full_path] = is_checked


# --- 主处理逻辑：现在它只依赖于 st.session_state.selected_files ---
img_files = st.session_state.get('selected_files', [])

if img_files:
    # ... (后续的识别、展示、聊天的所有代码都保持和上一版完全一样) ...
    @st.cache_resource
    def load_models(): return InsectDetector(), InsectIdentifier()
    detector, identifier = load_models()

    if len(img_files) > 1:
        st.info("已选择多张图片。识别结果将逐一展示，AI问答功能仅对第一张图片开启。")

    for i, (fname, image) in enumerate(img_files):
        st.markdown(f"---\n##### 图片：{fname}")
        if max(image.shape[:2]) > 1280: image = khandy.resize_image_long(image, 1280)
        image_for_draw = image.copy(); image_height, image_width = image.shape[:2]
        boxes, _, _ = detector.detect(image)
        cropped_list = []
        if not boxes.any():
            st.warning(f"图片 {fname} 未检测到昆虫，将对整张图片进行识别。")
            cropped_list = [(image, [0, 0, image_width, image_height])]
        else:
            box = boxes[0].astype(np.int32)
            if (box[2] - box[0] >= 30 and box[3] - box[1] >= 30):
                cropped_list.append((khandy.crop_or_pad(image, box[0], box[1], box[2], box[3]), box))
        
        if cropped_list:
            cropped, box = cropped_list[0]

            # all_probs = identifier.predict(cropped)[0]
            # top_5_indices = np.argsort(all_probs)[-5:][::-1]
            try:
                cropped_rgb = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB)
                pil_image = Image.fromarray(cropped_rgb)
                buffered = io.BytesIO()
                pil_image.save(buffered, format="JPEG")
                img_str_base64 = base64.b64encode(buffered.getvalue()).decode()
                
                # 调用ViT模型API进行推理
                vit_payload = {"image_base64": img_str_base64}
                vit_response = requests.post(VIT_API_URL, json=vit_payload, timeout=30)
                vit_response.raise_for_status()
                vit_result = vit_response.json()
                
                # 解析ViT模型的预测结果
                predictions = vit_result.get("predictions", [])
                if predictions:
                    # 构造类似原来格式的数据
                    sorted_predictions = sorted(predictions, key=lambda x: x["probability"], reverse=True)
                    top_5_predictions = sorted_predictions[:5]
                    table_data = []
            
                    for pred in top_5_predictions:
                        
                        table_data.append({
                            "类别 (中文名)": pred["class_name"].replace("_", " "),
                            "类别 (拉丁名)": "未知",  # 如果您有拉丁名可以填入
                            "概率": f"{pred['probability']:.3f}"
                        })
                    
                    all_probs = np.array([pred["probability"] for pred in predictions])
                    top_5_predictions = np.array([pred["class_index"] for pred in top_5_predictions])
                    
                    # 获取最可能得类别信息
                    top_1_pred = top_5_predictions[0]
                    top_1_prob = top_1_pred["probability"]
                    top_1_idx = top_1_pred["class_index"]
                    top_1_info = {
                        'chinese_name':top_1_pred["class_name"],
                        'latin_name': '未知'
                    }
                    # 构造文本显示
                    text = 'Unknown' if top_1_prob < 0.20 else f"{top_1_pred['class_name']}: {top_1_prob:.3f}"
                else:
                    # 没有预测结果的情况
                    all_probs = np.array([])
                    top_5_indices = np.array([])
                    table_data = [{"类别 (中文名)": "未知", "类别 (拉丁名)": "未知", "概率": "0.000"}]
                    top_1_prob = 0.0
                    top_1_idx = 0
                    top_1_info = {"chinese_name": "未知", "latin_name": "未知"}
                    text = "Unknown"
                    
                
            except Exception as e:
                st.error(f"调用ViT模型API时出错: {e}")
                # 出错时使用默认值
                all_probs = np.array([])
                top_5_indices = np.array([])
                table_data = [{"类别 (中文名)": "API错误", "类别 (拉丁名)": "API错误", "概率": "0.000"}]
                top_1_prob = 0.0
                top_1_idx = 0
                top_1_info = {"chinese_name": "未知", "latin_name": "未知"}
                text = "API Error"
                
            # for idx in top_5_indices:
            #     class_info = identifier.label_name_dict[idx]; prob = all_probs[idx]
            #     table_data.append({"类别 (中文名)": class_info['chinese_name'].replace('_', ' '), "类别 (拉丁名)": class_info['latin_name'], "概率": f"{prob:.3f}"})
            
            # top_1_idx = top_5_indices[0]; top_1_prob = all_probs[top_1_idx]; top_1_info = identifier.label_name_dict[top_1_idx]
            # text = 'Unknown' if top_1_prob < 0.20 else f"{top_1_info['chinese_name']}: {top_1_prob:.3f}"
            
            if len(boxes) > 0:
                position = [box[0] + 2, box[1] - 20]
                position[0] = min(max(position[0], 0), image_width); position[1] = min(max(position[1], 0), image_height)
                cv2.rectangle(image_for_draw, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
                try: image_for_draw = khandy.draw_text(image_for_draw, text, position, font='/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc', font_size=32)
                except Exception: cv2.putText(image_for_draw, text, (position[0], position[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            else:
                position = [10, 30]
                try: image_for_draw = khandy.draw_text(image_for_draw, text, position, font='/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc', font_size=32)
                except Exception: cv2.putText(image_for_draw, text, (position[0], position[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
            left_col, right_col = st.columns([1.2, 1.5], gap="large")
            with left_col:
                caption = "检测与识别结果" if len(boxes) > 0 else "原图（未检测到昆虫）"
                st.image(image_for_draw[..., ::-1], caption=caption, width=350)
            with right_col:
                st.subheader("Top 5 识别结果")
                st.markdown(
                """
                <style>
                    /* 调整表格字体大小 */
                    .stDataFrame table {
                        font-size: 20px; /* 设置合适的字体大小 */
                    }
                </style>
                """,
                unsafe_allow_html=True
                )
                df = pd.DataFrame(table_data)
                st.dataframe(df, use_container_width=True, height=210, hide_index=True)
            
            if i == 0:
                top_names = top_1_info['chinese_name'].split('_')
                mu, ke, zhong = '', '', ''
                for name in top_names:
                    if '目' in name: mu = name
                    elif '科' in name: ke = name
                if all(x not in top_names[-1] for x in ['目', '科', '属', '种']): zhong = top_names[-1]
                st.session_state["current_insect_info"] = {"image": image, "name": top_1_info['chinese_name'], "order": mu, "family": ke, "species": zhong, "latin_name": top_1_info['latin_name'], "probability": float(top_1_prob)}
        else:
            st.error(f"图片 {fname} 未能处理任何昆虫目标。")
            if i == 0: st.session_state.current_insect_info = None

    # --- 聊天界面UI和逻辑 (保持不变) ---
    if st.session_state.get("current_insect_info"):
        st.markdown("---")
        st.subheader("与AI聊聊 💬")
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]): st.markdown(message["content"])
        if prompt := st.chat_input("关于这个图像，有什么想问我的吗？"):
            st.session_state.chat_history.append({"role": "user", "content": prompt})
            with st.chat_message("user"): st.markdown(prompt)
            with st.chat_message("assistant"):
                response_generator = real_multimodal_chatbot(image_np=st.session_state["current_insect_info"]["image"], species_info=st.session_state["current_insect_info"], user_query=prompt, chat_history=st.session_state.chat_history[:-1])
                response = st.write_stream(response_generator)
            st.session_state.chat_history.append({"role": "assistant", "content": response})